{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 35,
   "source": [
    "import tensorflow as tf\n",
    "import scipy.optimize\n",
    "import scipy.io\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "from pyDOE import lhs         #Latin Hypercube Sampling\n",
    "import seaborn as sns \n",
    "import codecs, json\n",
    "\n",
    "tf.keras.backend.set_floatx('float64') # float64 as default\n",
    "\n",
    "# generates same random numbers each time\n",
    "np.random.seed(1234)\n",
    "tf.random.set_seed(1234)\n",
    "\n",
    "print(\"TensorFlow version: {}\".format(tf.__version__))"
   ],
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": [
      "TensorFlow version: 2.1.0\n"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# *Data Prep*\n",
    "\n",
    "Training and Testing data is prepared from the solution file"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "source": [
    "x_1 = np.linspace(-1,1,256)  # 256 points between -1 and 1 [256x1]\n",
    "x_2 = np.linspace(1,-1,256)  # 256 points between -1 and 1 [256x1]\n",
    "\n",
    "X, Y = np.meshgrid(x_1,x_2) "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Test Data\n",
    "\n",
    "We prepare the test data to compare against the solution produced by the PINN."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "source": [
    "X_u_test = np.hstack((X.flatten(order='F')[:,None], Y.flatten(order='F')[:,None]))\n",
    "\n",
    "# Domain bounds\n",
    "lb = np.array([-1, -1])\n",
    "ub = np.array([1, 1]) \n",
    "\n",
    "a_1 = 1\n",
    "a_2 = 1\n",
    "\n",
    "usol = np.sin(a_1 * np.pi * X) * np.sin(a_2 * np.pi * Y)\n",
    "\n",
    "u = usol.flatten('F')[:,None]"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Training Data"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "source": [
    "def trainingdata(N_u,N_f,N_v):\n",
    "    \n",
    "    leftedge_x = np.hstack((X[:,0][:,None], Y[:,0][:,None]))\n",
    "    leftedge_u = usol[:,0][:,None]\n",
    "    \n",
    "    rightedge_x = np.hstack((X[:,-1][:,None], Y[:,-1][:,None]))\n",
    "    rightedge_u = usol[:,-1][:,None]\n",
    "    \n",
    "    topedge_x = np.hstack((X[0,:][:,None], Y[0,:][:,None]))\n",
    "    topedge_u = usol[0,:][:,None]\n",
    "    \n",
    "    bottomedge_x = np.hstack((X[-1,:][:,None], Y[-1,:][:,None]))\n",
    "    bottomedge_u = usol[-1,:][:,None]\n",
    "    \n",
    "    all_X_u_train = np.vstack([leftedge_x, rightedge_x, bottomedge_x, topedge_x])\n",
    "    all_u_train = np.vstack([leftedge_u, rightedge_u, bottomedge_u, topedge_u])  \n",
    "     \n",
    "    #choose random N_u points for training\n",
    "    idx = np.random.choice(all_X_u_train.shape[0], N_u+N_v, replace=False) \n",
    "    \n",
    "    X_u_train = all_X_u_train[idx[0:N_u], :] #choose indices from  set 'idx' (x,t)\n",
    "    u_train = all_u_train[idx[0:N_u],:]      #choose corresponding u\n",
    "    \n",
    "    X_boundary_val = all_X_u_train[idx[N_u+1:N_u+N_v], :] #choose indices from  set 'idx' (x,t)\n",
    "    u_boundary_val = all_u_train[idx[N_u+1:N_u+N_v],:]      #choose corresponding u\n",
    "    \n",
    "    '''Collocation Points'''\n",
    "\n",
    "    # Latin Hypercube sampling for collocation points \n",
    "    # N_f sets of tuples(x,t)\n",
    "    X_f = lb + (ub-lb)*lhs(2,N_f+N_v) \n",
    "    X_f_train = np.vstack((X_f[0:N_f,:], X_u_train)) # append training points to collocation points \n",
    "    X_interior_val = X_f[N_f:N_f+N_v,:]\n",
    "    \n",
    "    return X_f_train, X_u_train, u_train, X_interior_val, X_boundary_val, u_boundary_val   \n"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# PINN "
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "source": [
    "def generate_PINN(n,L):  \n",
    "\n",
    "    #generate an array representing the network architecture\n",
    "    layers = np.ones(1+L).astype(int)*n # 1 input layer + L hidden layers\n",
    "    layers[0] = 2\n",
    "    layers[-1] = 1\n",
    "\n",
    "    #layers = [2, n1, n2, n3,........., n_L] #Network Configuration \n",
    "\n",
    "    PINN = tf.keras.Sequential()\n",
    "\n",
    "    PINN.add(tf.keras.layers.InputLayer(input_shape=(layers[0],),name=\"input_layer\", dtype = 'float64'))\n",
    "    \n",
    "    PINN.add(tf.keras.layers.Lambda(lambda X: 2*(X - lb)/(ub - lb) - 1)) \n",
    "    \n",
    "    initializer = 'glorot_normal'     \n",
    "    \n",
    "    for l in range (len(layers)-2):\n",
    "\n",
    "      # Xavier Initialization  \n",
    "      PINN.add(tf.keras.layers.Dense(layers[l+1],kernel_initializer=initializer, bias_initializer='zeros',\n",
    "                                  activation = tf.nn.tanh, name = \"layer\" + str(l+1), dtype = 'float64'))\n",
    "        \n",
    "    PINN.add(tf.keras.layers.Dense(layers[-1],kernel_initializer=initializer, bias_initializer='zeros',\n",
    "                          activation = None, name = \"output_layer\" , dtype = 'float64'))    \n",
    "        \n",
    "    return PINN  "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# *Loss Function*\n",
    "\n",
    "The loss function consists of two parts:\n",
    "1. **loss_BC**: MSE error of boundary losses\n",
    "2. **loss_PDE**: MSE error of collocation points satisfying the PDE\n",
    "\n",
    "**loss** = loss_BC + loss_PDE"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "source": [
    "def loss_BC(model, x, y):\n",
    "\n",
    "    loss_u = tf.reduce_mean(tf.square(y-model(x)))\n",
    "    return loss_u\n",
    "\n",
    "def loss_PDE(model, g):\n",
    "    \n",
    "    k = 1    \n",
    "        \n",
    "    x_1_f = g[:,0:1]\n",
    "    x_2_f = g[:,1:2]\n",
    "    \n",
    "    with tf.GradientTape(persistent=True) as tape:\n",
    "\n",
    "        tape.watch(x_1_f)\n",
    "        tape.watch(x_2_f)\n",
    "        \n",
    "        g = tf.stack([x_1_f[:,0], x_2_f[:,0]], axis=1)\n",
    "        \n",
    "        u = model(g)\n",
    "        u_x_1 = tape.gradient(u,x_1_f)\n",
    "        u_x_2 = tape.gradient(u,x_2_f)\n",
    "    \n",
    "    u_xx_1 = tape.gradient(u_x_1,x_1_f)\n",
    "    u_xx_2 = tape.gradient(u_x_2,x_2_f)\n",
    "    \n",
    "    del tape\n",
    "    \n",
    "    q = -( (a_1*np.pi)**2 + (a_2*np.pi)**2 - k**2 ) * np.sin(a_1*np.pi*x_1_f) * np.sin(a_2*np.pi*x_2_f)\n",
    "            \n",
    "    f = u_xx_1 + u_xx_2 + k**2 * u - q #residual\n",
    "    \n",
    "    loss_f = tf.reduce_mean(tf.square(f))\n",
    "    \n",
    "    return loss_f\n",
    "\n",
    "def loss(model,x,y,g):\n",
    "\n",
    "    loss_u = loss_BC(model,x,y)\n",
    "    loss_f = loss_PDE(model,g)\n",
    "\n",
    "    loss = loss_u + loss_f\n",
    "    \n",
    "    return loss\n",
    "\n",
    "def loss_split(model, x,y, g):\n",
    "\n",
    "    loss_u = loss_BC(model,x,y)\n",
    "    loss_f = loss_PDE(model,g)\n",
    "    \n",
    "    return loss_u,loss_f  "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# *Function Factory*\n",
    "\n",
    "A factory to create an optimization function and gradient as required by scipy.optimize"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "source": [
    "def function_factory(model, x, y, x_to_train_f):\n",
    "\n",
    "  \"\"\" \n",
    "    # Args:\n",
    "    #   model [in]: an instance of `tf.keras.Model` or its subclasses.\n",
    "    #   loss [in]: a function with signature loss_value = loss(pred_y, true_y).\n",
    "    #   x [in]: the input part of training data.\n",
    "    #   y [in]: the output part of training data.\n",
    "\n",
    "      Returns:\n",
    "        A function that has a signature of:\n",
    "            loss_value, gradients = f(model_parameters)\n",
    "  \"\"\"\n",
    "\n",
    "  #converting the input into a tensor for Automatic diff\n",
    "  g = tf.Variable(x_to_train_f, dtype = 'float64') \n",
    "\n",
    "  # obtain the shapes of all trainable parameters in the model\n",
    "  shapes = tf.shape_n(model.trainable_variables)\n",
    "  n_tensors = len(shapes)\n",
    "\n",
    "  # we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to\n",
    "  # prepare required information first\n",
    "  count = 0\n",
    "  idx = [] # stitch indices\n",
    "  part = [] # partition indices\n",
    "\n",
    "  for i, shape in enumerate(shapes):\n",
    "    n = np.product(shape)\n",
    "    idx.append(tf.reshape(tf.range(count, count+n, dtype=tf.int32), shape))\n",
    "    part.extend([i]*n)\n",
    "    count += n\n",
    "\n",
    "  part = tf.constant(part)\n",
    "\n",
    "  def assign_new_model_parameters(params_1d):\n",
    "      \"\"\"A function updating the model's parameters with a 1D tf.Tensor.\n",
    "\n",
    "      Args:\n",
    "          params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.\n",
    "      \"\"\"\n",
    "        \n",
    "      params = tf.dynamic_partition(params_1d, part, n_tensors)\n",
    "      for i, (shape, param) in enumerate(zip(shapes, params)):\n",
    "          model.trainable_variables[i].assign(tf.reshape(param, shape))\n",
    "\n",
    "  # now create a function that will be returned by this factory\n",
    "\n",
    "  def f(params_1d):\n",
    "    \"\"\"A function that can be used by an optimizer\n",
    "\n",
    "    This function is created by function_factory.\n",
    "\n",
    "    Args:\n",
    "      params_1d [in]: a 1D tf.Tensor.\n",
    "\n",
    "    Returns:\n",
    "        A scalar loss and the gradients w.r.t. the `params_1d`.\n",
    "    \"\"\"\n",
    "\n",
    "    with tf.GradientTape() as tape:\n",
    "      \n",
    "      assign_new_model_parameters(params_1d) # update the parameters in the model\n",
    "\n",
    "      loss_value = loss(model, x, y, g)      # calculate the loss\n",
    "\n",
    "    # calculate gradients and convert to 1D tf.Tensor\n",
    "    grads = tape.gradient(loss_value, model.trainable_variables)\n",
    "    grads = tf.dynamic_stitch(idx, grads)\n",
    "\n",
    "    return loss_value.numpy().astype(np.float64), grads.numpy().astype(np.float64)\n",
    "\n",
    "  # store this information as variables so we can use them outside the scope\n",
    "  f.iter = tf.Variable(0)\n",
    "  f.idx = idx\n",
    "  f.part = part\n",
    "  f.shapes = shapes\n",
    "  f.assign_new_model_parameters = assign_new_model_parameters\n",
    "      \n",
    "  return f"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "source": [
    "def metrics(networkparams):\n",
    "    \n",
    "    func.assign_new_model_parameters(networkparams)\n",
    "    \n",
    "    g = tf.Variable(X_f_train, dtype = 'float64')\n",
    "    \n",
    "    loss_value = loss(PINN, X_u_train, u_train, g)\n",
    "    \n",
    "    u_pred = PINN(X_u_test)\n",
    "    \n",
    "    error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)\n",
    "    \n",
    "    h = tf.Variable(X_interior_val, dtype = 'float64')\n",
    "        \n",
    "    val_loss = loss(PINN,X_boundary_val,u_boundary_val,h)\n",
    "    \n",
    "    text_file = open(\"loss_helm.txt\", \"a\")\n",
    "    n = text_file.write(str(loss_value.numpy())+ \" \" + str(error_vec) + \" \" + str(val_loss.numpy()) + \"\\n\" )\n",
    "    text_file.close()"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "# *Model Training and Testing*\n",
    "\n",
    "A function '**model**' is defined to generate a NN as per the input set of hyperparameters, which is then trained and tested. The L2 Norm of the solution error is returned as a comparison metric"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "N_u = N_v = 400 #Total number of data points for 'u'\n",
    "N_f = 100 #Total number of collocation points \n",
    "\n",
    "# Training data\n",
    "X_f_train, X_u_train, u_train, X_interior_val, X_boundary_val, u_boundary_val = trainingdata(N_u,N_f,N_v)\n",
    "\n",
    "n = 5  # Neurons in each hidden layer\n",
    "L = 4 # Number of hidden layers\n",
    "\n",
    "# Neural Network\n",
    "PINN = generate_PINN(n,L)\n",
    "\n",
    "func = function_factory(PINN, X_u_train, u_train, X_f_train)\n",
    "\n",
    "# convert initial model parameters to a 1D tf.Tensor\n",
    "init_params = tf.dynamic_stitch(func.idx, PINN.trainable_variables)\n",
    "\n",
    "start_time = time.time() \n",
    "\n",
    "# train the model with Scipy L-BFGS optimizer\n",
    "\n",
    "results = scipy.optimize.minimize(fun = func, \n",
    "                                  x0 = init_params, \n",
    "                                  args=(), \n",
    "                                  method='L-BFGS-B', \n",
    "                                  jac= True,        # If jac is True, fun is assumed to return the gradient along with the objective function\n",
    "                                  callback = None, \n",
    "                                  options = {'disp': None,\n",
    "                                            'maxcor': 200, \n",
    "                                            'ftol': 1 * np.finfo(float).eps,  #The iteration stops when (f^k - f^{k+1})/max{|f^k|,|f^{k+1}|,1} <= ftol\n",
    "                                            'gtol': 5e-5, \n",
    "                                            'maxfun':  50000, \n",
    "                                            'maxiter': 1,\n",
    "                                            'iprint': 50,   #print update every 50 iterations\n",
    "                                            'maxls': 50})\n",
    "\n",
    "\n",
    "elapsed = time.time() - start_time                \n",
    "print('Training time: %.2f' % (elapsed))\n",
    "\n",
    "# after training, the final optimized parameters are in results.position as vectors\n",
    "# so we have to manually put them back to the model as matrices\n",
    "func.assign_new_model_parameters(results.x)\n",
    "\n",
    "print(results)\n",
    "\n",
    "''' Model Accuracy ''' \n",
    "u_pred = PINN(X_u_test)\n",
    "\n",
    "error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)        # Relative L2 Norm of the error (Vector)\n",
    "print('Test Error: %.5f'  % (error_vec))\n",
    "\n",
    "u_pred = np.reshape(u_pred,(256,256),order='F') \n",
    "\n",
    "#Plot and save image\n",
    "plt.pcolor(x_1,x_2,u_pred, cmap = 'jet')\n",
    "plt.axis('scaled')\n",
    "# plt.savefig('Helmholtz_solution.png',dpi = 500)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "func.assign_new_model_parameters(results.x)\n",
    "g = tf.Variable(X_f_train, dtype = 'float64')\n",
    "\n",
    "shapes = tf.shape_n(PINN.trainable_variables) \n",
    "n_tensors = len(shapes)\n",
    "count = 0\n",
    "idx = [] # stitch indices\n",
    "part = [] # partition indices\n",
    "\n",
    "for i, shape in enumerate(shapes):\n",
    "    n = np.product(shape)\n",
    "    idx.append(tf.reshape(tf.range(count, count+n, dtype=tf.int32), shape))\n",
    "    part.extend([i]*n)\n",
    "    count += n\n",
    "    \n",
    "with tf.GradientTape(persistent=True) as tape_1:\n",
    "    \n",
    "    loss_u,loss_f = loss_split(PINN, X_u_train, u_train,g)\n",
    "\n",
    "    grad_u = tape_1.gradient(loss_u, PINN.trainable_variables) \n",
    "    grad_f = tape_1.gradient(loss_f, PINN.trainable_variables)\n",
    "    "
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "source": [
    "#Writing gradients to 2 separate JSON files \n",
    "for i in range (L*2):\n",
    "\n",
    "    temp = grad_f[i].numpy().tolist() # nested lists with same data, indices\n",
    "    json.dump(temp, codecs.open(\"Gradients/gradients_f\" + str(i) + \".json\", 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=0)\n",
    "    \n",
    "    temp = grad_u[i].numpy().tolist() # nested lists with same data, indices\n",
    "    json.dump(temp, codecs.open(\"Gradients/gradients_u\" + str(i) + \".json\", 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=0)\n",
    "    "
   ],
   "outputs": [
    {
     "output_type": "error",
     "ename": "NameError",
     "evalue": "name 'grad_f' is not defined",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-36-0a2f6644f223>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mL\u001b[0m\u001b[1;33m*\u001b[0m\u001b[1;36m2\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 4\u001b[1;33m     \u001b[0mtemp\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mgrad_f\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtolist\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;31m# nested lists with same data, indices\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      5\u001b[0m     \u001b[0mjson\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mdump\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mtemp\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mcodecs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mopen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"Gradients/gradients_f\"\u001b[0m \u001b[1;33m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mi\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m+\u001b[0m \u001b[1;34m\".json\"\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'w'\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mencoding\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;34m'utf-8'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mseparators\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m','\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m':'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0msort_keys\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;32mTrue\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mindent\u001b[0m\u001b[1;33m=\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      6\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'grad_f' is not defined"
     ]
    }
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "source": [
    "f = {}\n",
    "u = {}\n",
    "\n",
    "#Reading gradients from JSON files and storing in a dict \n",
    "for i in range (L*2):\n",
    "\n",
    "    obj_text = codecs.open(\"Gradients/gradients_f\" + str(i) + \".json\", 'r', encoding='utf-8').read()\n",
    "    f['f'+ str(i)] = np.array(json.loads(obj_text))\n",
    "    \n",
    "    obj_text = codecs.open(\"Gradients/gradients_u\" + str(i) + \".json\", 'r', encoding='utf-8').read()\n",
    "    u['u'+ str(i)] = np.array(json.loads(obj_text))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "sns.set_theme()\n",
    "for i in range (2*L):\n",
    "    \n",
    "    d = {}\n",
    "    \n",
    "    fig = plt.figure()\n",
    "    \n",
    "    d['grad_f'] = f['f'+ str(i)].flatten()\n",
    "    d['grad_u'] = u['u'+ str(i)].flatten()\n",
    "    fig = sns.displot(d,kind=\"kde\").savefig(\"Gradients/output\" + str(i) + \".png\", dpi=500)\n",
    "    \n",
    "#     fig.savefig(\"Gradients/output\" + str(i) + \".png\", dpi=500)\n"
   ],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}